#!/usr/bin/env python3
import os
import argparse
import logging
import shutil
import sys

def process_scene(scene, processed_root, depth_root, merge):
    """
    Check correspondence of .JPG → _depth.png for a single scene.
    If all files match, merge the processed scene dir into the depth dir.
    Returns True on full correspondence, False otherwise.
    """
    proc_scene_dir = os.path.join(processed_root, scene)
    depth_scene_dir = os.path.join(depth_root, scene)

    # Make sure it's a directory
    if not os.path.isdir(proc_scene_dir):
        logging.warning(f"Skipping {scene!r}: not a directory")
        return False

    # Collect all .JPG/.jpg files
    jpg_files = [f for f in os.listdir(proc_scene_dir)
                 if f.lower().endswith('.jpg')]

    if not jpg_files:
        logging.info(f"No JPG files found in scene {scene!r}, skipping.")
        return False

    missing_files = {}
    # Check each jpg has a corresponding depth file
    for jpg in jpg_files:
        base = jpg.split('_')[0]
        depth_fname = f"{base}_depth.png"
        depth_fpath = os.path.join(depth_scene_dir, depth_fname)
        if not os.path.exists(depth_fpath):
            missing_files[jpg] = depth_fname
            break
            # return False   
    if missing_files:
        logging.warning(f"Missing depth files for {scene!r}: {missing_files}")
        return False
    
    logging.info(f"Found all depth files for {scene!r}")
    # All matched: merge processed RGB files into depth scene folder
    # Just like in Ubuntu: mv -f proc_scene_dir depth_scene_dir
    try:
        if merge:
            # Ensure depth_scene_dir exists
            os.makedirs(depth_scene_dir, exist_ok=True)
            # Copy all files from proc_scene_dir to depth_scene_dir with overwrite
            shutil.copytree(proc_scene_dir, depth_scene_dir, dirs_exist_ok=True)
            # Remove the processed scene directory
            shutil.rmtree(proc_scene_dir)

            logging.info(f"Merged {scene!r} into depth directory")
    except Exception as e:
        logging.error(f"Failed to merge {scene!r}: {str(e)}")
        return False

    return True

def main():
    parser = argparse.ArgumentParser(
        description="Check and merge ScanNet++ processed RGBs with their depth maps"
    )
    parser.add_argument("--processed_dir", type=str, default="/lc/data/3D/scannetpp_processed", help="Root of scannetpp_processed")
    parser.add_argument("--depth_dir", type=str, default="/lc/data/3D/scannetpp_depth/nomask", help="Root of scannetpp_depth/nomask")
    parser.add_argument("--log_file", type=str, default="/lc/data/3D/scannetpp_depth/merge.txt",  help="Path to write log")
    parser.add_argument(
    "--merge",
    dest="merge",
    action="store_true",
    help="Enable merging")
    args = parser.parse_args()

    # Set up logging
    logging.basicConfig(
        filename=args.log_file,
        level=logging.INFO,
        format="%(message)s",
        datefmt="%Y-%m-%d %H:%M:%S"
    )

    if not os.path.isdir(args.processed_dir):
        logging.critical(f"Processed directory not found: {args.processed_dir}")
        sys.exit(1)
    if not os.path.isdir(args.depth_dir):
        logging.critical(f"Depth directory not found: {args.depth_dir}")
        sys.exit(1)

    missing = []
    success = []
    # Iterate over all scene subdirectories
    scenes = sorted(os.listdir(args.processed_dir))
    for scene in scenes:
        if process_scene(scene, args.processed_dir, args.depth_dir, args.merge):
            success.append(scene)
        else:
            missing.append(scene)

    logging.info(f"Processed {len(success)} scenes: {success}")
    logging.info(f"Failed to process {len(missing)} scenes: {missing}")
    logging.info("All scenes processed.")

if __name__ == "__main__":
    main()
